import numpy as np

class RectangularMesh:
    """矩形计算网格类"""
    def __init__(self, nx, ny):
        self.nx = nx
        self.ny = ny
        # 根据实际尺寸(42x69cm)设置网格间距
        self.dx = 0.42 / nx  # 米
        self.dy = 0.69 / ny  # 米

class CFDSolver:
    def __init__(self, mesh):
        self.mesh = mesh
        self.u = np.zeros((mesh.nx+1, mesh.ny+2))  # x方向速度场
        self.v = np.zeros((mesh.nx+2, mesh.ny+1))  # y方向速度场
        self.p = np.zeros((mesh.nx+2, mesh.ny+2))  # 压力场
        self.T = np.ones((mesh.nx+2, mesh.ny+2)) * 298.0  # 温度场初始化为环境温度
        self.max_val = 1.0e10  # 设置数值上限，防止溢出

        # 定义几何参数（米）
        self.cpu_width = 0.09  # 9cm
        self.cpu_height = 0.07  # 7cm
        self.cpu_bottom = 0.60  # 距顶部9cm
        self.cpu_left = 0.06  # 距左边6cm
        self.cpu_right = 0.27  # 距右边6cm

        self.gpu_width = 0.02  # 2cm
        self.gpu_height = 0.28  # 28cm
        self.gpu_spacing = 0.002  # 2mm
        self.gpu_bottom = 0.04  # 距底部4cm

        # 左侧矩形参数
        self.left_rect_width = 0.10  # 10cm
        self.left_rect_height = 0.35  # 35cm

    def set_boundary_conditions(self, inlet_vel, inlet_temp, gpu_power, gpu_positions):
        """设置边界条件"""
        nx, ny = self.mesh.nx, self.mesh.ny
        dx, dy = self.mesh.dx, self.mesh.dy

        # 重置场
        self.u.fill(0.0)
        self.v.fill(0.0)
        self.T.fill(inlet_temp)

        # 1. 入口边界（顶部）- 冷气入口
        self.u[:, -1] = inlet_vel
        self.T[:, -1] = inlet_temp

        # 2. CPU散热器设置
        # 左侧CPU
        i_start = int(self.cpu_left / dx)
        i_end = int((self.cpu_left + self.cpu_width) / dx)
        j_start = int((self.cpu_bottom - self.cpu_height) / dy)
        j_end = int(self.cpu_bottom / dy)

        # CPU底边设置为热空气出口
        self.T[i_start:i_end, j_start] = 343.15  # 70°C
        self.v[i_start:i_end, j_start] = 0.2  # 0.2m/s

        # 右侧CPU
        i_start = int(self.cpu_right / dx)
        i_end = int((self.cpu_right + self.cpu_width) / dx)
        self.T[i_start:i_end, j_start] = 343.15  # 70°C
        self.v[i_start:i_end, j_start] = 0.2  # 0.2m/s

        # 3. GPU设置
        for pos in gpu_positions:
            # 计算GPU位置
            x_start = int((0.10 + 0.036 + pos * (self.gpu_width + self.gpu_spacing)) / dx)
            x_end = int(x_start + self.gpu_width / dx)
            y_start = int(self.gpu_bottom / dy)
            y_end = int((self.gpu_bottom + self.gpu_height) / dy)

            # GPU上、下、右边界设置为热空气出口
            self.T[x_start:x_end, y_start:y_end] = 348.15  # 75°C

            # 设置出口速度
            self.u[x_end, y_start:y_end] = 0.2  # 右边界
            self.v[x_start:x_end, y_start] = 0.2  # 下边界
            self.v[x_start:x_end, y_end] = 0.2  # 上边界

    def setup_grid(self):
        """设置网格"""
        # 设置网格大小
        self.nx = 52  # x方向网格数
        self.ny = 52  # y方向网格数

        # 计算网格间距
        self.dx = self.Lx / (self.nx - 1)
        self.dy = self.Ly / (self.ny - 1)

        # 创建网格点坐标
        self.x = np.linspace(0, self.Lx, self.nx)
        self.y = np.linspace(0, self.Ly, self.ny)

        # 初始化速度场、压力场和温度场
        self.u = np.zeros((self.nx, self.ny))  # x方向速度
        self.v = np.zeros((self.nx, self.ny))  # y方向速度
        self.p = np.zeros((self.nx, self.ny))  # 压力
        self.T = np.full((self.nx, self.ny), self.T_ambient)  # 温度场初始化为环境温度

        # 标记电源区域（左下角10x35cm区域）
        self.power_supply_mask = np.zeros((self.nx, self.ny), dtype=bool)
        x_power = int(0.10 / self.Lx * (self.nx - 1))  # 10cm对应的网格点
        y_power = int(0.35 / self.Ly * (self.ny - 1))  # 35cm对应的网格点

        for i in range(self.nx):
            for j in range(self.ny):
                if i <= x_power and j <= y_power:
                    self.power_supply_mask[i,j] = True

        # 在电源区域设置边界条件
        self.u[self.power_supply_mask] = 0
        self.v[self.power_supply_mask] = 0
        self.T[self.power_supply_mask] = self.T_ambient  # 电源区域不参与温度计算

    def solve_momentum_x(self):
        """求解x方向动量方程"""
        u_new = self.u.copy()

        for i in range(1, self.nx-1):
            for j in range(1, self.ny-1):
                if not self.power_supply_mask[i,j]:  # 跳过电源区域
                    # ... 原有的动量方程求解代码 ...
                    u_new[i,j] = self.u[i,j] - self.dt * (
                        self.u[i,j] * (self.u[i+1,j] - self.u[i-1,j]) / (2*self.dx) +
                        self.v[i,j] * (self.u[i,j+1] - self.u[i,j-1]) / (2*self.dy)
                    ) - self.dt / self.rho * (self.p[i+1,j] - self.p[i-1,j]) / (2*self.dx)

        return u_new

    def solve_momentum_y(self):
        """求解y方向动量方程"""
        v_new = self.v.copy()

        for i in range(1, self.nx-1):
            for j in range(1, self.ny-1):
                if not self.power_supply_mask[i,j]:  # 跳过电源区域
                    # ... 原有的动量方程求解代码 ...
                    v_new[i,j] = self.v[i,j] - self.dt * (
                        self.u[i,j] * (self.v[i+1,j] - self.v[i-1,j]) / (2*self.dx) +
                        self.v[i,j] * (self.v[i,j+1] - self.v[i,j-1]) / (2*self.dy)
                    ) - self.dt / self.rho * (self.p[i,j+1] - self.p[i,j-1]) / (2*self.dy)

        return v_new

    def solve_energy(self):
        """求解能量方程"""
        T_new = self.T.copy()

        for i in range(1, self.nx-1):
            for j in range(1, self.ny-1):
                if not self.power_supply_mask[i,j]:  # 跳过电源区域
                    # ... 原有的能量方程求解代码 ...
                    T_new[i,j] = self.T[i,j] - self.dt * (
                        self.u[i,j] * (self.T[i+1,j] - self.T[i-1,j]) / (2*self.dx) +
                        self.v[i,j] * (self.T[i,j+1] - self.T[i,j-1]) / (2*self.dy)
                    )

        return T_new

    def solve(self, algorithm="SIMPLE"):
        """主求解方法"""
        print(algorithm)
        if algorithm == "SIMPLE":
            return self._solve_SIMPLE()
        elif algorithm == "PISO":
            return self._solve_PISO()
        elif algorithm == "SIMPLEC":
            return self._solve_SIMPLEC()
        else:
            print(f"未知算法: {algorithm}，使用默认SIMPLE算法")
            return self._solve_SIMPLE()

    def _solve_SIMPLE(self, max_iter=10, tol=1e-5):
        """SIMPLE算法实现"""
        residual_history = []

        # 初始化边界条件
        self._apply_boundary_conditions()

        for iter in range(max_iter):
            print("迭代次数：", iter)
            # 动量方程求解
            self._solve_momentum_equations()

            # 压力修正
            p_corr = self._pressure_correction()

            # 速度修正
            self._correct_velocity(p_corr)

            # 更新压力场
            self.p += 0.7 * p_corr  # 松弛因子

            # 限制压力范围
            self.p = np.clip(self.p, -self.max_val, self.max_val)

            # 能量方程求解（温度场）
            self._solve_energy_equation()

            # 收敛判断
            res = np.max(np.abs(p_corr))
            print("残差：", res)
            if not np.isfinite(res):  # 检查是否为有效值
                print("警告: 计算发散，停止迭代")
                break

            residual_history.append(res)
            if res < tol:
                break

        # 计算最终温度场
        self._calculate_final_temperature()

        print("求解结束：", residual_history)
        return residual_history

    def _solve_PISO(self, max_iter=10, tol=1e-5):
        """PISO算法实现"""
        residual_history = []

        # 初始化边界条件
        self._apply_boundary_conditions()

        for iter in range(max_iter):
            print("迭代次数：", iter)

            # 动量方程求解
            self._solve_momentum_equations()

            # 第一次压力修正
            p_corr1 = self._pressure_correction()

            # 第一次速度修正
            self._correct_velocity(p_corr1)

            # 第二次压力修正
            p_corr2 = self._pressure_correction()

            # 第二次速度修正
            self._correct_velocity(p_corr2)

            # 更新压力场 (PISO不需要松弛)
            self.p += p_corr2

            # 限制压力范围
            self.p = np.clip(self.p, -self.max_val, self.max_val)

            # 能量方程求解
            self._solve_energy_equation()

            # 收敛判断
            res = np.max(np.abs(p_corr2))
            print("残差：", res)
            if not np.isfinite(res):  # 检查是否为有效值
                print("警告: 计算发散，停止迭代")
                break

            residual_history.append(res)
            if res < tol:
                break

        # 计算最终温度场
        self._calculate_final_temperature()
        return residual_history

    def _solve_SIMPLEC(self, max_iter=10, tol=1e-5):
        """SIMPLEC算法实现"""
        residual_history = []

        # 初始化边界条件
        self._apply_boundary_conditions()

        for iter in range(max_iter):
            print("迭代次数：", iter)

            # 动量方程求解
            self._solve_momentum_equations()

            # 压力修正 (SIMPLEC使用一致性修正)
            p_corr = self._pressure_correction_consistent()

            # 速度修正 (SIMPLEC不使用松弛)
            self._correct_velocity_consistent(p_corr)

            # 更新压力场 (SIMPLEC使用更高的松弛因子)
            self.p += 0.9 * p_corr

            # 限制压力范围
            self.p = np.clip(self.p, -self.max_val, self.max_val)

            # 能量方程求解
            self._solve_energy_equation()

            # 收敛判断
            res = np.max(np.abs(p_corr))
            print("残差：", res)
            if not np.isfinite(res):  # 检查是否为有效值
                print("警告: 计算发散，停止迭代")
                break

            residual_history.append(res)
            if res < tol:
                break

        # 计算最终温度场
        self._calculate_final_temperature()
        return residual_history

    def _apply_boundary_conditions(self):
        """应用边界条件，确保初始场的稳定性"""
        nx, ny = self.mesh.nx, self.mesh.ny

        # 确保温度场有合理值
        self.T = np.clip(self.T, 273.15, 373.15)  # 限制温度在0°C到100°C之间

        # 设置边界条件
        # 底部出口
        self.T[:, 0] = self.T[:, 1]
        self.v[:, 0] = 0.2  # 出口速度

        # 左边界（壁面）
        self.T[0, :] = self.T[1, :]
        self.u[0, :] = 0.0  # 无滑移边界

        # 右边界（壁面）
        self.T[-1, :] = self.T[-2, :]
        self.u[-1, :] = 0.0  # 无滑移边界

        # 确保速度场有合理值
        self.u = np.clip(self.u, -10.0, 10.0)
        self.v = np.clip(self.v, -10.0, 10.0)

    def _solve_momentum_equations(self):
        """动量方程求解"""
        nx, ny = self.mesh.nx, self.mesh.ny
        dx, dy = self.mesh.dx, self.mesh.dy

        # 简化的隐式求解
        u_new = np.copy(self.u)
        v_new = np.copy(self.v)

        for i in range(1, nx):
            for j in range(1, ny+1):
                # x方向动量方程
                u_new[i, j] = 0.5 * (self.u[i+1, j] + self.u[i-1, j] + self.u[i, j+1] + self.u[i, j-1]) - \
                              0.5 * dx * (self.p[i+1, j] - self.p[i, j])

        for i in range(1, nx+1):
            for j in range(1, ny):
                # y方向动量方程
                v_new[i, j] = 0.5 * (self.v[i+1, j] + self.v[i-1, j] + self.v[i, j+1] + self.v[i, j-1]) - \
                              0.5 * dy * (self.p[i, j+1] - self.p[i, j])

                # 考虑浮力效应 (温度影响)
                try:
                    v_new[i, j] += 0.01 * (self.T[i, j] - 298.0)  # 简化的浮力模型
                except:
                    # 如果出现异常，保持原值
                    pass

        # 限制速度范围，防止发散
        u_new = np.clip(u_new, -10.0, 10.0)
        v_new = np.clip(v_new, -10.0, 10.0)

        # 更新速度场
        self.u = u_new
        self.v = v_new

    def _pressure_correction(self):
        """压力修正方程"""
        nx, ny = self.mesh.nx, self.mesh.ny
        p_corr = np.zeros_like(self.p)

        # 简化的Poisson方程求解
        for k in range(50):  # 迭代求解
            p_old = np.copy(p_corr)

            for i in range(1, nx+1):
                for j in range(1, ny+1):
                    try:
                        div = (self.u[i, j] - self.u[i-1, j]) + (self.v[i, j] - self.v[i, j-1])
                        p_corr[i, j] = 0.25 * (p_corr[i+1, j] + p_corr[i-1, j] +
                                              p_corr[i, j+1] + p_corr[i, j-1] - div)
                    except:
                        # 如果出现异常，使用前一个值
                        p_corr[i, j] = p_old[i, j]

            # 限制修正值范围
            p_corr = np.clip(p_corr, -self.max_val, self.max_val)

            # 检查收敛性
            if np.max(np.abs(p_corr - p_old)) < 1e-5:
                break

        return p_corr

    def _pressure_correction_consistent(self):
        """SIMPLEC一致性压力修正方程"""
        nx, ny = self.mesh.nx, self.mesh.ny
        p_corr = np.zeros_like(self.p)

        # SIMPLEC修正系数
        alpha = 0.95

        # 简化的Poisson方程求解 (一致性修正)
        for k in range(50):  # 迭代求解
            p_old = np.copy(p_corr)

            for i in range(1, nx+1):
                for j in range(1, ny+1):
                    try:
                        div = (self.u[i, j] - self.u[i-1, j]) + (self.v[i, j] - self.v[i, j-1])
                        p_corr[i, j] = 0.25 * alpha * (p_corr[i+1, j] + p_corr[i-1, j] +
                                                      p_corr[i, j+1] + p_corr[i, j-1] - div)
                    except:
                        # 如果出现异常，使用前一个值
                        p_corr[i, j] = p_old[i, j]

            # 限制修正值范围
            p_corr = np.clip(p_corr, -self.max_val, self.max_val)

            # 检查收敛性
            if np.max(np.abs(p_corr - p_old)) < 1e-5:
                break

        return p_corr

    def _correct_velocity(self, p_corr):
        """速度修正"""
        nx, ny = self.mesh.nx, self.mesh.ny
        dx, dy = self.mesh.dx, self.mesh.dy

        # 修正x方向速度
        for i in range(1, nx):
            for j in range(1, ny+1):
                try:
                    self.u[i, j] -= 0.5 * (p_corr[i+1, j] - p_corr[i, j]) / dx
                except:
                    pass  # 忽略异常

        # 修正y方向速度
        for i in range(1, nx+1):
            for j in range(1, ny):
                try:
                    self.v[i, j] -= 0.5 * (p_corr[i, j+1] - p_corr[i, j]) / dy
                except:
                    pass  # 忽略异常

        # 限制速度范围
        self.u = np.clip(self.u, -10.0, 10.0)
        self.v = np.clip(self.v, -10.0, 10.0)

    def _correct_velocity_consistent(self, p_corr):
        """SIMPLEC一致性速度修正"""
        nx, ny = self.mesh.nx, self.mesh.ny
        dx, dy = self.mesh.dx, self.mesh.dy

        # 修正x方向速度 (一致性修正)
        for i in range(1, nx):
            for j in range(1, ny+1):
                try:
                    self.u[i, j] -= 0.8 * (p_corr[i+1, j] - p_corr[i, j]) / dx
                except:
                    pass  # 忽略异常

        # 修正y方向速度 (一致性修正)
        for i in range(1, nx+1):
            for j in range(1, ny):
                try:
                    self.v[i, j] -= 0.8 * (p_corr[i, j+1] - p_corr[i, j]) / dy
                except:
                    pass  # 忽略异常

        # 限制速度范围
        self.u = np.clip(self.u, -10.0, 10.0)
        self.v = np.clip(self.v, -10.0, 10.0)

    def _solve_energy_equation(self):
        """能量方程求解 (温度场)"""
        nx, ny = self.mesh.nx, self.mesh.ny

        # 简化的能量方程求解
        T_new = np.copy(self.T)

        for i in range(1, nx+1):
            for j in range(1, ny+1):
                try:
                    # 对流-扩散方程
                    advection = (self.u[i, j] * (self.T[i+1, j] - self.T[i-1, j]) +
                                self.v[i, j] * (self.T[i, j+1] - self.T[i, j-1])) * 0.25

                    diffusion = 0.1 * (self.T[i+1, j] + self.T[i-1, j] +
                                      self.T[i, j+1] + self.T[i, j-1] - 4*self.T[i, j])

                    T_new[i, j] = self.T[i, j] + 0.1 * (diffusion - advection)
                except:
                    # 如果计算失败，保持原温度
                    T_new[i, j] = self.T[i, j]

        # 限制温度范围，防止非物理值
        T_new = np.clip(T_new, 273.0, 373.0)  # 限制温度在0°C到100°C之间

        # 更新温度场
        self.T = T_new

    def _calculate_final_temperature(self):
        """计算最终温度分布"""
        # 应用边界条件和热源
        nx, ny = self.mesh.nx, self.mesh.ny

        # 确保边界条件正确
        self.T[:, 0] = 2 * self.T[:, 1] - self.T[:, 2]  # 底部出口
        self.T[0, :] = self.T[1, :]  # 左边界
        self.T[nx+1, :] = self.T[nx, :]  # 右边界

        # 限制温度范围，确保物理合理性
        self.T = np.clip(self.T, 273.0, 373.0)  # 限制温度在0°C到100°C之间